import torch.optim as optim
import inspect

# 定义所有PyTorch常用优化器及其支持的参数映射
OPTIMIZER_PARAMS = {
    'SGD': {
        'required': [],
        'optional': ['lr', 'momentum', 'dampening', 'weight_decay', 'nesterov', 'maximize', 'foreach', 'differentiable'],
        'description': '随机梯度下降优化器'
    },
    'Adam': {
        'required': [],
        'optional': ['lr', 'betas', 'eps', 'weight_decay', 'amsgrad', 'maximize', 'foreach', 'capturable', 'differentiable', 'fused'],
        'description': 'Adam自适应学习率优化器'
    },
    'Adamax': {
        'required': [],
        'optional': ['lr', 'betas', 'eps', 'weight_decay', 'maximize', 'foreach'],
        'description': '基于无穷范数的Adam变体'
    },
    'RMSprop': {
        'required': [],
        'optional': ['lr', 'alpha', 'eps', 'weight_decay', 'momentum', 'centered', 'maximize', 'foreach'],
        'description': 'RMSprop优化器'
    },
    'Adagrad': {
        'required': [],
        'optional': ['lr', 'lr_decay', 'weight_decay', 'initial_accumulator_value', 'eps', 'maximize', 'foreach'],
        'description': 'Adagrad优化器'
    },
    'Adadelta': {
        'required': [],
        'optional': ['lr', 'rho', 'eps', 'weight_decay', 'maximize', 'foreach'],
        'description': 'Adadelta优化器'
    },
    'AdamW': {
        'required': [],
        'optional': ['lr', 'betas', 'eps', 'weight_decay', 'amsgrad', 'maximize', 'foreach', 'capturable', 'differentiable', 'fused'],
        'description': '带权重衰减的Adam优化器'
    },
    'SparseAdam': {
        'required': [],
        'optional': ['lr', 'betas', 'eps'],
        'description': '适用于稀疏张量的Adam变体'
    },
    'LBFGS': {
        'required': [],
        'optional': ['lr', 'max_iter', 'max_eval', 'tolerance_grad', 'tolerance_change', 'history_size', 'line_search_fn'],
        'description': 'L-BFGS优化器'
    }
}


def get_optimizer_supported_params(optimizer_type):
    """
    获取指定优化器支持的参数列表
    
    参数:
        optimizer_type (str): 优化器类型名称
    
    返回:
        dict: 包含required和optional参数的字典
    """
    if optimizer_type in OPTIMIZER_PARAMS:
        return OPTIMIZER_PARAMS[optimizer_type]
    return {'required': [], 'optional': []}


def filter_optimizer_params(optimizer_class, input_params):
    """
    智能过滤优化器参数，只保留优化器支持的参数
    
    参数:
        optimizer_class: 优化器类
        input_params (dict): 输入的参数字典
    
    返回:
        dict: 过滤后的参数字典
    """
    # 创建参数副本以避免修改原始数据
    params_copy = input_params.copy()
    
    try:
        # 使用inspect模块获取优化器构造函数的签名
        signature = inspect.signature(optimizer_class.__init__)
        valid_params = list(signature.parameters.keys())
        
        # 移除self参数
        if 'self' in valid_params:
            valid_params.remove('self')
        # 移除params参数
        if 'params' in valid_params:
            valid_params.remove('params')
        
        # 过滤参数，只保留有效的参数
        filtered_params = {}
        removed_params = []
        
        for key, value in params_copy.items():
            if key in valid_params:
                # 特殊处理某些参数类型转换
                if key == 'eps' and isinstance(value, str):
                    filtered_params[key] = float(value)
                elif key == 'betas' and isinstance(value, list):
                    filtered_params[key] = tuple(value)
                else:
                    filtered_params[key] = value
            else:
                removed_params.append(key)
        
        if removed_params:
            print(f"优化器 {optimizer_class.__name__} 不支持以下参数，已自动过滤: {removed_params}")
        
        return filtered_params
    except Exception as e:
        print(f"[警告] 无法自动检测优化器参数，使用备选方案: {str(e)}")
        # 备选方案：使用预定义的参数映射
        optimizer_type = optimizer_class.__name__
        supported = get_optimizer_supported_params(optimizer_type)
        
        if supported:
            valid_params = supported['required'] + supported['optional']
            filtered_params = {k: v for k, v in params_copy.items() if k in valid_params}
            removed_params = [k for k in params_copy.keys() if k not in valid_params]
            
            if removed_params:
                print(f"优化器 {optimizer_type} 不支持以下参数，已自动过滤: {removed_params}")
            
            return filtered_params
        
        # 如果都失败了，返回原始参数（可能会出错，但至少提供警告）
        print(f"[警告] 无法获取优化器 {optimizer_type} 的参数信息，将尝试使用所有提供的参数")
        return params_copy


def create_optimizer(model, optimizer_config):
    """
    根据配置动态创建优化器，自动处理参数兼容性
    
    参数:
        model: PyTorch模型实例
        optimizer_config (dict): 优化器配置字典，包含'type'和'args'字段
    
    返回:
        torch.optim.Optimizer: 创建的优化器实例
    
    异常:
        ValueError: 当优化器不存在或创建失败时抛出
    """
    if not isinstance(optimizer_config, dict) or 'type' not in optimizer_config:
        raise ValueError("优化器配置必须包含'type'字段")
    
    optimizer_type = optimizer_config['type']
    optimizer_args = optimizer_config.get('args', {})
    
    # 检查优化器是否存在
    if not hasattr(optim, optimizer_type):
        raise ValueError(f"PyTorch中不存在优化器: {optimizer_type}")
    
    try:
        optimizer_class = getattr(optim, optimizer_type)
        
        # 过滤参数字典
        filtered_args = filter_optimizer_params(optimizer_class, optimizer_args)
        
        # 创建优化器实例
        optimizer = optimizer_class(model.parameters(), **filtered_args)
        
        print(f"成功创建优化器: {optimizer_type}")
        print(f"使用的参数: {filtered_args}")
        
        return optimizer
    except Exception as e:
        raise ValueError(f"创建优化器 {optimizer_type} 失败: {str(e)}")


def get_optimizer_info(optimizer_type):
    """
    获取优化器的详细信息
    
    参数:
        optimizer_type (str): 优化器类型名称
    
    返回:
        dict: 包含优化器信息的字典
    """
    info = {
        'name': optimizer_type,
        'available': hasattr(optim, optimizer_type)
    }
    
    if optimizer_type in OPTIMIZER_PARAMS:
        info.update(OPTIMIZER_PARAMS[optimizer_type])
    
    return info


# 使用示例
def example_usage():
    """
    提供使用示例
    """
    # 示例1：基本使用
    """
    from training.utils.optimizer_util import create_optimizer
    
    # 假设我们有一个模型
    # model = YourModel()
    
    # 优化器配置
    optimizer_config = {
        'type': 'AdamW',
        'args': {
            'lr': 0.001,
            'weight_decay': 0.01,
            'eps': 1e-8,
            'betas': [0.9, 0.999],
            'momentum': 0.9  # 这个参数AdamW不支持，但会被自动过滤
        }
    }
    
    # 创建优化器
    # optimizer = create_optimizer(model, optimizer_config)
    """
    
    # 示例2：在配置文件中使用
    """
    # config.yaml示例
    """
    # 示例3：处理不同优化器
    """
    # 对于SGD
    sgd_config = {
        'type': 'SGD',
        'args': {
            'lr': 0.01,
            'momentum': 0.9,
            'weight_decay': 0.0001
        }
    }
    
    # 对于Adam
    adam_config = {
        'type': 'Adam',
        'args': {
            'lr': 0.001,
            'betas': [0.9, 0.999],
            'eps': 1e-8
        }
    }
    
    # 对于RMSprop
    rmsprop_config = {
        'type': 'RMSprop',
        'args': {
            'lr': 0.01,
            'alpha': 0.99,
            'eps': 1e-8,
            'momentum': 0.9
            }
        }
    """